import streamlit as st
from create_plan9 import create_plan


# 设置页面配置
st.set_page_config(layout="wide")

st.header("内运排班系统")
st.write("请输入以下信息，用于生成今日排班表")

#初始化任务
if not st.session_state:
    st.session_state = {
        '铝卷' : {'count': None,
                  'discard_drop_count': 1,
                  'discard_count' : 1},
        '铝棒' : {'count': None, 'load_times': []},
        '单顶' : {'count': None, 'load_times': []},
        '双零箔': {
            'count': None,
            'four_point_two': None,
            'ten_point_five': None,
            'thirteen': None,
            'load_time': '7:50'

        },
        '废铝卷':{
            'count': None,
            'scania': None,
            'benz': None
        },
        '长沙甩挂人员': {'name': None},
        '常德甩挂人员': {'name': None},
        '休息人员名单': {'rest_names':[]},
        '休息车辆': {'rest_car_numbers':[]}
    }

left_col, right_col = st.columns(2)
with left_col:

    st.subheader("铝卷信息")
    col1, col2, col3 =  st.columns(3)
    with col1 :
        st.session_state['铝卷']['count'] = st.number_input(
            '铝卷数量',min_value = 0,key = 'aluminum_roll_count'
        )
    with col2 :
        st.session_state['铝卷']['discard_drop_count'] = st.number_input(
            '铝卷(甩废铝)数量',min_value = 0,key = 'discard_drop_count'
        )
    with col3 :
        st.session_state['铝卷']['discard_count'] = st.number_input(
            '铝卷(废铝)数量',min_value = 0,key = 'discard_count'
        )

    st.subheader("铝棒信息")
    aluminum_rod_count =st.number_input(
            '铝棒数量',min_value = 0,key = 'aluminum_rod_count'
        )

    if aluminum_rod_count != st.session_state['铝棒']['count']:
        st.session_state['铝棒']['count'] = aluminum_rod_count
        st.session_state['铝棒']['load_times'] = [None] * aluminum_rod_count
        st.rerun()

    # 根据铝棒数量显示时间输入框
    if st.session_state['铝棒']['count'] > 0:
        st.subheader("为每趟铝棒输入装货时间")
        # 创建列来组织输入框
        cols = st.columns(2)

    for i in range(aluminum_rod_count):
        with cols[i%2]:
            # 使用唯一的key
            time_key = f"aluminum_rod_time_{i}"
            # 获取当前值或使用None
            current_value = st.session_state['铝棒']['load_times'][i] if i < len(
                st.session_state['铝棒']['load_times']) else None

            # 时间输入框
            selected_time = st.time_input(
                f"第 {i + 1} 趟铝棒装货时间",
                value=current_value,
                step = 1800,
                key=time_key
            )

            # 更新session state中的时间
            st.session_state['铝棒']['load_times'][i] = selected_time

    st.subheader("双零箔信息")
    st.session_state['双零箔']['count'] = st.number_input(
        '双零箔数量', min_value=0, key='double_zero_count'
    )
    col1, col2, col3 = st.columns(3)
    with col1:
        st.session_state['双零箔']['four_point_two'] = st.number_input(
            '4.2米车辆数量', min_value=0, key='four_point_two'
        )
    with col2:
        st.session_state['双零箔']['ten_point_five'] = st.number_input(
            '10.5米车辆数量', min_value=0, key='ten_point_five'
        )
    with col3:
        st.session_state['双零箔']['thirteen'] = st.number_input(
            '13米车辆数量', min_value=0, key='thirteen'
        )

    st.subheader('休息人员名单')
    rest_names_input = st.text_input(
        "请输入休息人员姓名，多个姓名用逗号分隔，例如“丁冬，杨军”", key="rest_names"
    )
    if rest_names_input:
        st.session_state['休息人员名单']['rest_names'] = [name.strip() for name in rest_names_input.split("，")]
    else:
        st.session_state['休息人员名单']['rest_names'] = []

with right_col:
    st.subheader("废铝卷信息")
    col1, col2, col3 =  st.columns(3)
    with col1 :
        st.session_state['废铝卷']['count'] = st.number_input(
            '废铝卷数量',min_value = 0,key = 'discard_aluminum_roll_count'
        )
    with col2 :
        st.session_state['废铝卷']['scania'] = st.number_input(
            '斯堪尼亚车数量',min_value = 0,key = 'scania_count'
        )
    with col3 :
        st.session_state['废铝卷']['benz'] = st.number_input(
            '奔驰车数量',min_value = 0,key = 'benz_count'
        )

    st.subheader("单顶信息")
    single_top_count = st.number_input(
        '单顶数量', min_value=0, key='single_top_count'
    )

    if single_top_count != st.session_state['单顶']['count']:
        st.session_state['单顶']['count'] = single_top_count
        st.session_state['单顶']['load_times'] = [None] * single_top_count
        st.rerun()

    # 根据单顶数量显示时间输入框
    if st.session_state['单顶']['count'] > 0:
        st.subheader("为每趟单顶输入装货时间")
        # 创建列来组织输入框
        cols = st.columns(2)

    for i in range(single_top_count):
        with cols[i % 2]:
            # 使用唯一的key
            time_key = f"single_top_time_{i}"
            # 获取当前值或使用None
            current_value = st.session_state['单顶']['load_times'][i] if i < len(
                st.session_state['单顶']['load_times']) else None

            # 时间输入框
            selected_time = st.time_input(
                f"第 {i + 1} 趟单顶装货时间",
                value=current_value,
                step=1800,
                key=time_key
            )

            # 更新session state中的时间
            st.session_state['单顶']['load_times'][i] = selected_time

    st.subheader("甩挂人员信息")
    st.session_state['长沙甩挂人员']['name'] =st.text_input(
            '长沙甩挂人员', key = 'changsha_drop_roll_name'
        )
    st.session_state['常德甩挂人员']['name'] = st.text_input(
        '常德甩挂人员', key='changde_drop_roll_name'
    )

    st.subheader('休息车辆名单')
    rest_car_numbers_input = st.text_input(
            "请输入休息车辆的车牌号，多个车牌号用逗号分隔，例如”湘AM9210，湘AK1377“", key="rest_car_numbers"
        )
    if rest_car_numbers_input:
        st.session_state['休息车辆']['rest_car_numbers'] = [number.strip() for number in rest_car_numbers_input.split("，")]
    else:
        st.session_state['休息车辆']['rest_car_numbers'] = []


if st.button('确认生成计划', type='primary'):
    try:
        result_df = create_plan(today_tasks=st.session_state)

        for col in result_df.columns:
            if result_df[col].dtype == 'object':
                result_df[col] = result_df[col].astype(str)

        st.success('计划生成成功！')
        st.session_state['df'] = result_df.copy()

    except Exception as e:
        st.error(f'生成计划失败，请重新输入任务信息: {e}')

# 始终显示数据编辑器（如果数据存在）
if 'df' in st.session_state and not st.session_state['df'].empty:
    # 使用 data_editor 并自动更新 session_state
    st.session_state['df'] = st.data_editor(
        st.session_state['df'],
        num_rows="dynamic",
        width='stretch',
        height=400,
        key='main_editor'
    )



